from openbabel import pybel
from openbabel import openbabel
import os
import shutil


def collect_error_sdfs(input_filename: str,
                       mol_id: str,
                       directory: str = 'splited_3d') -> None:
    '''
    Split the given sdf file without other changes.

    Then rename them by an unique sdf property,say datebase id.

    Then try to convert molecules and collect bad molecules' raw sdf files.
    '''

    # ------mkdirs------
    os.mkdir(directory)
    os.mkdir(f'{directory}/good')
    os.mkdir(f'{directory}/bad')

    # ------split-------
    os.system(f'obabel {input_filename} -ocopy -O {directory}/split.sdf -m')
    error_logger = openbabel.obErrorLog
    error_logger.SetOutputLevel(openbabel.obError)  #只显示错误信息
    num_errors = 0
    for file in os.listdir(directory):
        if not os.path.isfile(f'{directory}/{file}'):
            continue
        mol = next(pybel.readfile('sdf', f'{directory}/{file}'))
        title = mol.data[mol_id]
        titled_filename = f'{directory}/{title}.sdf'
        os.rename(f'{directory}/{file}', titled_filename)

        # ----classify------
        tmp_filename = f'{titled_filename[:-4]}_tmp.sdf'
        with pybel.Outputfile('sdf', tmp_filename) as tmp_file:
            mol.title = title
            mol.addh()  #加氢
            mol.make3D()  #生成3D坐标
            tmp_file.write(mol)
        num_err_msgs = error_logger.GetErrorMessageCount()
        if num_err_msgs > num_errors:
            num_errors = num_err_msgs
            shutil.move(titled_filename, f'{directory}/bad/')
            print(f'{mol_id}:{title} meets error.')
        else:
            shutil.move(tmp_filename, f'{directory}/good/{title}.sdf')


if __name__ == '__main__':
    collect_error_sdfs('a.sdf', 'ID')
